#!/usr/bin/env python

# This script uses framestats logfiles to plot the distribution of I and P frames for a number of videos
#
# Inputs(given as arguments to script call): video1_framestats.txt, video2_framestats.txt etc...
# Output: Plot #2
#
# To do:
#	Check that x axis is correctly scaled (It seems so but messy code regarding kilobyte or bytes)
#	Implement verifying calculations (?) check if dist's sum to 1 etc...

import numpy as np
import pylab as pl
import scipy.stats as spstats
import os.path as osp
import sys
from matplotlib.font_manager import FontProperties as FP

def new_frame(filep):
	line = '1'
	while not line.startswith('[/FRAME]'):
		line = filep.readline()
		if line.startswith('pkt_pos='):
			pos = int(line[line.find('=')+1:])
		if line.startswith('pict_type='):
			pict_type = line[line.find('=')+1:-1]
	return pos, pict_type


# Feed this function with a filepointer and a reference to the plot, it then takes care of plotting. Will be called multiple times with the different logfiles.

def plot(filep,video_nb,tracker):
	
	pos_type = []

	##### DATA EXTRACTION #####

	line = '1'
	while line:
		line = f.readline()
		if line.startswith('[FRAME]'):
			position, pict_type = new_frame(f)
			pos_type.append((position, pict_type))

	frame_sizes = []
	p_size_array = []
	i_size_array = []

	for i in range(len(pos_type)-1):
		size = pos_type[i+1][0] - pos_type[i][0]
		size = size/1000.
		frame_sizes.append(size)
		if str(pos_type[i][1]).lower() == 'p':
			p_size_array.append(size)
		else:
			i_size_array.append(size)

	##### Data preparation #######

	# Video stats
	psize = np.sum(p_size_array) # Size of all p frames
	isize = np.sum(i_size_array) # Size of all i frames
	totalsize = isize+psize # total video size

	## NB: x-axis was normalized before - not anymore

#	# Normalize i_size_array and p_size_array before making histogram
#	n_i_size_array=[]
#	n_p_size_array=[]

#	for k in range(len(i_size_array)):
#		n_i_size_array.append(i_size_array[k]/float(totalsize))
#		
#	for k in range(len(p_size_array)):
#		n_p_size_array.append(p_size_array[k]/float(totalsize))

	# We are no longer normalizing
	n_i_size_array=i_size_array
	n_p_size_array=p_size_array

	# Not normalizing
	max_size=100 # [kB]
	min_size=0

	# Check if we choose max_size properly
	if max_size<max(i_size_array) or max_size<max(p_size_array):
		print 'WARNING: Not all frames are plotted, increase max_size [kB]'
		
	histogram_bins_i = 50 # Choose wisely !
	histogram_bins_p = 50 # Choose wisely !

	hist_i = pl.histogram(n_i_size_array, histogram_bins_i, (min_size, max_size))
	hist_p = pl.histogram(n_p_size_array, histogram_bins_p, (min_size, max_size))

	# Print useful information to terminal
	print 'Video '+str(video_nb+1)
	print 'Total filesize '+str(totalsize/1000.)+' MB'
	print 'I-frame data percentage', (isize/float(totalsize))*100
	print 'I-frame average size', isize/len(i_size_array)
	print 'P-frame data percentage', (psize/totalsize)*100
	print 'P-frame average size', psize/len(p_size_array)
	p_size_mb = psize/1000.
	p_size_mb_str = str("%.1f" % p_size_mb)
	i_size_mb = isize/1000.
	i_size_mb_str = str("%.1f" % i_size_mb)
	p_label = 'P-frames ('+str(len(p_size_array))+' frames, '+p_size_mb_str+' MB)'
	print p_label
	i_label = 'I-frames ('+str(len(i_size_array))+' frames, '+i_size_mb_str+' MB)'
	print i_label

	# (P) Normalize y data to percentage of P frames respectively
	norm_hist_p_y=[]

	for k in range(len(hist_p[0])):
		norm_hist_p_y.append(hist_p[0][k]/float(len(p_size_array)))
		
	x_data_p=hist_p[1][:-1]
	y_data_p=norm_hist_p_y
	
	# (P) Normalize y data to percentage of P frames respectively
	norm_hist_i_y=[]

	for k in range(len(hist_i[0])):
		norm_hist_i_y.append(hist_i[0][k]/float(len(i_size_array)))
		
	x_data_i=hist_i[1][:-1]
	y_data_i=norm_hist_i_y

	# Separate plots for I and P
	fig = pl.figure(figsize=[10,4]) # Setting aspect ratio
	fig.subplots_adjust(wspace=0.5) # spacing between subplots
	ax1 = fig.add_subplot(121) # left plot
	ax2 = fig.add_subplot(122,sharex=ax1,sharey=ax1) # right plot, share axes

	# Layout fix: Convert 0 to -1 in y_data which cleans plot
	for k in range(len(y_data_p)):
		if y_data_p[k]==0:
			y_data_p[k]=-1

	for k in range(len(y_data_i)):
		if y_data_i[k]==0:
			y_data_i[k]=-1

	# Layout fix: Numerate legend with letters: A,B,C (expect to only go to C) ...
	# How about .1 .2 .3 ?

	if video_nb<3:
		letter='A.'+str((tracker%3)+1)
	elif video_nb<6:
		letter='B.'+str((tracker%3)+1)
	else:
		letter='C.'+str((tracker%3)+1)
		
	# Update plot with data
#	ax1.plot(x_data_p,y_data_p,linestyle='-',label='Video '+str(video_nb+1)+' P-frames',drawstyle='steps-post')
#	ax2.plot(x_data_i,y_data_i,linestyle='-',label='Video '+str(video_nb+1)+' I-frames',drawstyle='steps-post')

	ax1.plot(x_data_p,y_data_p,linestyle='-',color='black',label='Video '+letter+' P-frames',drawstyle='steps-post')
	ax2.plot(x_data_i,y_data_i,linestyle='-',color='black',label='Video '+letter+' I-frames',drawstyle='steps-post')

	mode_value_p_index=np.argmax(y_data_p)
	mode_value_p=x_data_p[mode_value_p_index]
	print 'P-frame mode value is '+ str(mode_value_p)+' [kB]'

#	mode_p_frame_data=spstats.mode(y_data_p)
#	print 'P-frame mode value '+str(mode_p_frame_data[0]*100)+' kB'

	# Plot annotation
#	fig.autofmt_xdate() too much marketing
	ax1.set_ylabel('Part of all P-frames [-]')
	ax1.set_xlabel('Frame size [kB]')
	ax1.legend(prop=fp)
	ax1.grid(True,which='both')
	ax2.set_ylabel('Part of all I-frames [-]')
	ax2.set_xlabel('Frame size [kB]')
	ax2.legend(prop=fp)
	ax2.grid(True,which='both')
	pl.xlim(0,max_size)
	pl.ylim(0,0.40)

	# Save figure
	workingdir = osp.abspath('../../plotting/video_coding2')
	pl.savefig(workingdir+'/figs/'+'dist_frames_'+str(video_nb+1)+'.eps')
	
if __name__ == '__main__':

	# Font fix
	fp = FP()
	fp.set_size('small')

	# tracker for legend notation
	tracker=0

	# Run through all of the given framestats.txt logfiles given in arguments
	for arg in range(len(sys.argv)-1):
		input_filename = sys.argv[arg+1]
		f = open(input_filename,'r')
		plot(f,arg,tracker)
		tracker+=1

	sys.exit()




